import sys
import numpy as np
from tqdm import tqdm
import argparse

if '../../../embeddings' not in sys.path:
    sys.path.append('../../../embeddings')

from seq2tensor import s2t
from tensorflow.keras.models import Model
from tensorflow.keras.layers import (
    Dense, Bidirectional, concatenate, multiply, LeakyReLU,
    Conv1D, MaxPooling1D, GlobalAveragePooling1D, Input, GRU
)

def build_model(hidden_dim=25, seq_size=600, dim=15):
    seq_input1 = Input(shape=(seq_size, dim), name='seq1')
    seq_input2 = Input(shape=(seq_size, dim), name='seq2')
    l1=Conv1D(hidden_dim, 3)
    r1=Bidirectional(GRU(hidden_dim, return_sequences=True))
    l2=Conv1D(hidden_dim, 3)
    r2=Bidirectional(GRU(hidden_dim, return_sequences=True))
    l3=Conv1D(hidden_dim, 3)
    r3=Bidirectional(GRU(hidden_dim, return_sequences=True))
    l4=Conv1D(hidden_dim, 3)
    r4=Bidirectional(GRU(hidden_dim, return_sequences=True))
    l5=Conv1D(hidden_dim, 3)
    r5=Bidirectional(GRU(hidden_dim, return_sequences=True))
    l6=Conv1D(hidden_dim, 3)
    s1=MaxPooling1D(3)(l1(seq_input1))
    s1=concatenate([r1(s1), s1])
    s1=MaxPooling1D(3)(l2(s1))
    s1=concatenate([r2(s1), s1])
    s1=MaxPooling1D(2)(l3(s1))
    s1=concatenate([r3(s1), s1])
    s1=MaxPooling1D(2)(l4(s1))
    s1=concatenate([r4(s1), s1])
    s1=MaxPooling1D(2)(l5(s1))
    s1=concatenate([r5(s1), s1])
    s1=l6(s1)
    s1=GlobalAveragePooling1D()(s1)
    s2=MaxPooling1D(3)(l1(seq_input2))
    s2=concatenate([r1(s2), s2])
    s2=MaxPooling1D(3)(l2(s2))
    s2=concatenate([r2(s2), s2])
    s2=MaxPooling1D(2)(l3(s2))
    s2=concatenate([r3(s2), s2])
    s2=MaxPooling1D(2)(l4(s2))
    s2=concatenate([r4(s2), s2])
    s2=MaxPooling1D(2)(l5(s2))
    s2=concatenate([r5(s2), s2])
    s2=l6(s2)
    s2=GlobalAveragePooling1D()(s2)
    merge_text = multiply([s1, s2])
    x = Dense(hidden_dim, activation='linear')(merge_text)
    x = LeakyReLU(alpha=0.3)(x)
    x = Dense(int((hidden_dim+7)/2), activation='linear')(x)
    x = LeakyReLU(alpha=0.3)(x)
    main_output = Dense(2, activation='softmax')(x)
    merge_model = Model(inputs=[seq_input1, seq_input2], outputs=[main_output])
    return merge_model


# 参数解析
parser = argparse.ArgumentParser()
parser.add_argument("--model_path", type=str, required=True, help="训练好的模型权重路径")
parser.add_argument("--emb_file", type=str, required=True, help="使用的嵌入文件路径")
parser.add_argument("--input_pairs", type=str, required=True, help="待预测的蛋白对文件路径")
parser.add_argument("--output", type=str, default="predictions.tsv", help="预测结果输出路径")
args = parser.parse_args()

# 初始化序列嵌入系统
seq2t = s2t(args.emb_file)
seq_size = 600

# 加载模型
model = build_model()
print(model.summary())
model.load_weights(args.model_path)

def process_sequences(seq_pairs):
    """处理蛋白质序列对为模型输入格式"""
    tensor_list1, tensor_list2 = [], []
    for seq1, seq2 in seq_pairs:
        tensor1 = seq2t.embed_normalized(seq1, seq_size)
        tensor2 = seq2t.embed_normalized(seq2, seq_size)
        tensor_list1.append(tensor1)
        tensor_list2.append(tensor2)
    return np.array(tensor_list1), np.array(tensor_list2)

# 读取输入文件（假设格式：seq1\tseq2）
test_pairs = []
with open(args.input_pairs) as f:
    for line in f:
        seq1, seq2 = line.strip().split('\t')
        test_pairs.append((seq1, seq2))

# 生成预测输入
input1, input2 = process_sequences(test_pairs)

# 执行预测
predictions = model.predict([input1, input2])

# 输出结果
with open(args.output, 'w') as fw:
    fw.write("Sequence1\tSequence2\tProbability\tPrediction\n")
    for (seq1, seq2), pred in zip(test_pairs, predictions):
        prob = pred[0]  # 假设索引0对应结合概率
        cls = 1 if prob > 0.5 else 0
        fw.write(f"{seq1}\t{seq2}\t{prob:.4f}\t{cls}\n")
